#! /usr/bin/env python
"""
Convert one or more still image sequences into QuickTime movies.

Author: Chip Chapin <cchapin@gmail.com>
For More Info: http://code.google.com/p/seq-to-qt/

You must have the PyWin32 COM interface installed,
see http://sourceforge.net/projects/pywin32/

Acknowledgements:
  The QuickTime interface was originally written as a JScript+WSH
  script, which I based on a 2006 blog post by Luc-Eric Rousseau
  (XSIBlog http://www.xsi-blog.com/archives/103).  Rousseau's script
  was in turn based on sample code by John Cromie, author of the
  book "QuickTime for .NET and COM Developers" (Elsevier 2006,
  http://www.skylark.ie/qt4.net/samplecode.asp).

Latest Update: 2010-02-08

TODO: Some sort of progress indicator during the initial rendering stage.
TODO: Better solution to the QT_PLAYER_DELAY hack.
"""
import datetime
import math
import os
import re
import shutil
import sys
import tempfile
import time
import win32com.client

# Codec info file will be saved in your working directory.
# Saving it in "C:\\..." can fail if you lack permission.
CODEC_INFO_FILENAME = "qtMovieFromStillsCodecInfo.xml"
DEFAULT_CHUNK_FILES=7000
DEFAULT_FRAME_RATE = "60"

# HACK: Delay (sec) while the QuickTime player initializes.
QT_PLAYER_DELAY = 7.0

def _abort(message=None):
    """Print error message and exit."""
    if message:
        print "ERROR: %s" % message
    sys.exit(1)

def _log(message):
    print time.strftime("%Y-%m-%d %X  " + message)

def _usage(message=None):
    """Print usage message and exit."""
    if message:
        print message
    sys.exit(2)

def check_file_folder(fpath):
    """Returns true if the directory of fpath exists and is writable."""
    fdir = os.path.dirname(fpath)
    try:
        f = tempfile.TemporaryFile("w", dir=fdir)
    except OSError:
        return False
    f.close()
    return True

def file_exists(fpath):
    """Returns true if file fpath exists and is readable."""
    try:
        f = open(fpath, "r")
    except IOError:
        return False
    f.close()
    return True

def uniqueify(fpath):
    """Check to see if the file exists.  If so, uniqueify the filename."""
    upath = fpath
    count = 0
    while file_exists(upath):
        count = count + 1
        (root, ext) = os.path.splitext(fpath)
        upath = "%s-%02d%s" % (root, count, ext)
    if count:
        print "WARN: You already have a file '%s'" % fpath
        print "   Saving as '%s' instead." % upath
    return upath

def create_new_movie_from_images(sourcePath, frameRate, qtControl):
    """Create the movie from the still image sequence."""
    _log("Creating new movie from still sequence '%s' ..." % sourcePath)
    try:
        qtControl.CreateNewMovieFromImages(sourcePath,
                                           frameRate,
                                           True) # rate is in frames per second
    except:
        _log("ERROR creating movie ")
        #+ e.number +
        #           " (" + (e.number>>16 & 0x1FFF) +
        #           "-" + (e.number & 0xffff) + ")");
        #WScript.Echo(e.description);
        # TODO: Find more reliable error reporting.
        if not qtControl:
            print "Lost QuickTime connection"
        else:
            print "QuickTime error %d" % qtControl.ErrorCode
            qte = qtControl.QuickTime.Error
            print "  %d, %s " % (qte.ErrorCode, qte.Description)
            print "  %s" % qte.SourceReference
        raise

    if not qtControl.Movie:
        # Possibly instead of aborting we should skip to the next sequence.
        _abort("No movie created (" + qtControl.ErrorCode + ")")
    qtMovie = qtControl.Movie
    duration = qtMovie.Duration
    if (duration == 0):
        # This test isn't as helpful as I thought it would be.  I thought it
        # would catch the case where QT does not have a valid input file, but
        # in that case QT seems to create a two-second empty movie. 
        _abort("Movie has duration 0.")

    # Duration is the number of frames * 10.
    _log("Created new movie, duration %d" % duration)

def get_quicktime_exporter(qtControl):
    """Set up the QuickTime movie exporter."""
    qt = qtControl.QuickTime
    if (qt.Exporters.Count == 0):
        # Only add an exporter if needed.
        qt.Exporters.Add()
        print "Adding new Exporter."
    else:
        print "Using existing Exporter."
    qtExporter = qt.Exporters(1)
    if not qtExporter:
        _abort("Unable to get Exporter.")
    qtExporter.TypeName = "QuickTime Movie"
    qtExporter.SetDataSource(qtControl.Movie)

    if file_exists(CODEC_INFO_FILENAME):
        print "Reading codec config from '%s'" % CODEC_INFO_FILENAME
        CodecFileInfo = open(CODEC_INFO_FILENAME, "r")
        xmlCodecInfoText = CodecFileInfo.read()
        # Cause the exporter to be reconfigured.
        # http://developer.apple.com/technotes/tn2006/tn2120.html
        tempSettings = qtExporter.Settings
        tempSettings.XML = xmlCodecInfoText
        qtExporter.Settings = tempSettings
    else:
        # Use the Settings dialog box, then save the results.
        qtExporter.ShowSettingsDialog()
        xmlCodecInfoText = qtExporter.Settings.XML
        try:
            CodecFileInfo = open(CODEC_INFO_FILENAME, "w")
            CodecFileInfo.write(xmlCodecInfoText)
            CodecFileInfo.close()
        except IOError:
            print ("Warning: failed to save codec info to '%s'" %
                   CODEC_INFO_FILENAME)
            print "continuing ..."
    return qtExporter

def export_movie(qtExporter, destPath):
    """Export the movie."""
    _log("Exporting ...")
    try:
        qtExporter.DestinationFileName = destPath
        qtExporter.ShowProgressDialog = True
        # Uncomment this line if you want the export dialog box to appear.
        # qtExporter.ShowExportDialog();
        qtExporter.BeginExport()  # This can take a l-o-n-g time.
        _log("Exported to '%s'" % destPath)
    except:
        _log("ERROR exporting '%s'" % destPath)
        #print "ERROR " + (e.number>>16 & 0x1FFF) +
        #      "-" + (e.number & 0xffff) + 
        #      " exporting '" + destPath + "'"
        #print e.description
        #WScript.Echo(JSON.stringify(e, null, 2));
        #qte = qt.Error;
        #print "QuickTime Error %d, %s" % (qte.ErrorCode, qte.Description)
        #WScript.Echo(JSON.stringify(qte, null, 2));
        raise

def qt_movie_from_stills(sourcePath, destPath, frameRate):
    """Create and export a QuickTime movie from a sequence of images."""
    # Launch QuickTime Player Application
    qtPlayerApp = win32com.client.Dispatch(
                    "QuickTimePlayerLib.QuickTimePlayerApp")
    time.sleep(QT_PLAYER_DELAY);  # Give it time to launch.
    if not qtPlayerApp:
        _abort("Failed to launch QuickTime Player App.")

    # Get the QuickTime player and its associated controller.
    # NOTE: The script will abort here if the player hasn't had time
    # to initialize.  It should work if you run it again, or you can increase
    # QT_PLAYER_DELAY.
    qtPlayer = qtPlayerApp.Players(1);
    if not qtPlayer:
        _abort("Failed to get QuickTime Player.");
    print "Got player '%s'" % qtPlayer.Caption
    qtControl = qtPlayer.QTControl
    
    create_new_movie_from_images(sourcePath, frameRate, qtControl)
    qtExporter = get_quicktime_exporter(qtControl)
    export_movie(qtExporter, destPath)
    # Closing the player causes failures for subsequent invocations.
    # qtPlayer.Close();

def seq_file_list(a_seq_file):
    """Determine the list of file names in a sequence.
    
    Returns a duple (sequence_list, sequence_name)
    """
    # Extract the leading invariant part of the sequence filename.
    # For example, if the name is Fooo-0001.jpg, the invariant is "Fooo-".
    # Use non-greedy match so the numbers are kept out.
    (work_folder, sequence_proto) = os.path.split(a_seq_file)
    sm = re.match("(.*?)[0-9]+\.", sequence_proto)
    if not sm:
        _usage(("'%s' doesn't look like the start of a sequence."
              % sequence_proto) +
"\nSequences have sequentially numbered file names like Foo-001.jpg, Foo-002.jpg...\n" +
"  sequenceExample -- Path to any file in the sequence.\n" +
"  maxSize -- Maximum number of sequence files to put in a folder.");
        
    # Construct a regexp for matching sequence file names.
    sq_name = sm.group(1)
    sq_re = re.compile("^" + sq_name + "([0-9]+)\.")

    # Make a list of files that match the sq_re pattern.  "fnmatch" patterns
    # are not powerful enough to be reliable so we don't use glob.
    sq_files = []
    dir = os.listdir(work_folder)
    # See issue http://code.google.com/p/seq-to-qt/issues/detail?id=1
    # A lexical "dir.sort()" here can go quite wrong.  We do our own
    # "smart sort" on the numeric sequence numbers.
    # TODO: Find a more elegant way to do this.
    for f in dir:
        sqfm = sq_re.match(f)
        if sqfm:
            sq_files.append((int(sqfm.group(1)), f))
    sq_files.sort()
    sqf = []
    for (num, f) in sq_files:
        sqf.append(f)
    return (sqf, sq_name)

def split_file_list(a_seq_file, sq_files, sq_name, max_size):
    """Split the sequence files into successive split folders.
    Returns a list of (folder, starting_file) duples.
    """
    work_folder = os.path.dirname(a_seq_file)

    sq_count = len(sq_files)
    num_splits = math.ceil(float(sq_count) / max_size)
    print ("Splitting file sequence '%s' into %d parts."
           % (sq_name, num_splits))
    parent_folder = os.path.dirname(work_folder)
    file_count = 0
    next_file_in_split = 1
    current_split = 0
    splits = []
    for fname in sq_files:
        if (file_count == 0 or next_file_in_split > max_size):
            current_split = current_split + 1
            next_file_in_split = 1
            split_folder = os.path.join(parent_folder,
                                        ("%s-%d" % (work_folder, current_split)))
            _log("New split folder '%s' (%d)" % (split_folder, file_count))
            try:
                os.mkdir(split_folder, 0755)
            except EnvironmentError as (errno, strerror):
                _log("ERROR: Failed to create split directory '%s'" %
                     split_folder)
                print "   [Error %d] %s" % (errno, strerror)
                _abort()
            splits.append((split_folder, fname))
            
        # Move the file from work_folder to split_folder
        # print fname + " => " + split_folder
        shutil.move(os.path.join(work_folder, fname),
                    os.path.join(split_folder, fname))
        file_count = file_count + 1
        next_file_in_split = next_file_in_split + 1

    _log("Split %d files in '%s' into %d folders." %
         (file_count, work_folder, current_split))
    return splits

def make_seq_chunks(a_seq_file, max_size):
    """Split a file sequence into folders of no more than max_size files each.
    Returns a list of (folder, starting_file) duples.

    Sequences have file names that end with a sequence number,
    like (Foo-001.jpg, Foo-002.jpg, ...).
    """
    (sq_files, sq_name) = seq_file_list(a_seq_file)
    sq_count = len(sq_files)
    print "Sequence '%s' contains %d members." % (a_seq_file, sq_count)    
    if sq_count <= max_size:
        print "No need to split."
        if (sq_count == 0):
            return []
        else:
            return [os.path.split(a_seq_file)]
    else:
        return split_file_list(a_seq_file, sq_files, sq_name, max_size)

def seq_name_from_filename(fname):
    """Extract the 'sequence name' from a representative filename."""
    sm = re.match("(.*?)[-_]?[0-9]+\.", os.path.basename(fname))
    if not sm:
        print ("'%s' doesn't look like the start of a sequence."
                  % fname)
        return None
    return sm.group(1)

def convert_sequences(base_dir, seqs,
                      frame_rate=DEFAULT_FRAME_RATE,
                      chunk_size=DEFAULT_CHUNK_FILES):
    """Generates QuickTime movies for a set of image sequences.

    base_dir: Full path to the common parent directory.
    seqs: List of duples describing image sequences.
        Each duple contains the FOLDER NAME of the sequence
        and the STARTING IMAGE file name: (folder, file)
        Names are unqualified, i.e. relative to their parent directories.
    frameRate: Numeric frame rate in frames per second.
    chunk_size: Maximum number of frames to be processed in
        a single movie.
    """
    for (sq_dir, sq_file) in seqs:
        sq_src = os.path.join(base_dir, sq_dir, sq_file)
        if not file_exists(sq_src):
            print "WARN: '%s' is not readable.  Skipping." % sq_src
            continue
        # Possibly split the image sequence into chunks.
        splits = make_seq_chunks(sq_src, chunk_size)
        # Iterate over the chunks.
        for (split_folder, split_file) in splits:
            sq_name = seq_name_from_filename(sq_file)
            if (sq_name):
                src = os.path.join(split_folder, split_file)
                dstfile = ("%s_%s.mov" %
                           (sq_name, os.path.basename(split_folder)))
                dst = os.path.join(base_dir, dstfile)
                dst = uniqueify(dst)
                _log("Converting:") 
                print ("    Folder %s\n    Start %s\n    Dest %s"
                       % (split_folder, split_file, dst))
                qt_movie_from_stills(src, dst, frame_rate)
                _log("Complete.")

def main(argv):
    """Main Program"""
    _usage("This is the seq_to_qt module.  Use the commands instead:\n" +
"  make_seq_chunks.py -- Split a large sequence into multiple directories\n" +
"  qt_movie_from_stills.py -- Create a movie from an image sequence\n" +
"  seq_batch_convert.py -- Template for batch conversion.")

if __name__ == "__main__":
    main(sys.argv)
